import os
import fitz  # PyMuPDF
import pytesseract
from PIL import Image
import io
import threading

global_store = {}
def search_in_pdf(file_path, search_term):
    if '作废' in file_path:
        return
    doc = fitz.open(file_path)
    for page_num, page in enumerate(doc):
        text = page.get_text("text")
        if not text:
            image_list = page.get_images(full=True)
            for image_index, img in enumerate(image_list):
                # 获取图片的 XREF
                xref = img[0]
                # 提取图像字节
                base_image = doc.extract_image(xref)
                image_bytes = base_image["image"]
                # 读取图像
                image = Image.open(io.BytesIO(image_bytes))
                # 使用 PyTesseract 识别图像中的文字
                text = text + pytesseract.image_to_string(image, lang='chi_sim')
        global_store[f'{file_path}_{page_num + 1}'] = text
        if search_term.lower() in text.lower():
            print(f"Found '{search_term}' in {file_path} on page {page_num + 1}")
    doc.close()

def search_in_folder(folder_path, search_term):
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.endswith(".pdf"):
                search_in_pdf(os.path.join(root, file), search_term)


pytesseract.pytesseract.tesseract_cmd = r'E:\D其他软件\提效工具\tesseract-ocr\tesseract.exe'
search_word = "国家"
# 文件夹目录，"./" 表示当前目录
# folder_path = "./03_规章制度/"
folder_path = "./TestImagePDF/"
# search_in_folder(folder_path, search_word)
if __name__ == '__main__':
    while True:
        search_word = input("请输入你要查找的内容，然后点击回车^_^\n")
        if len(global_store) != 0:
            for key, value in global_store.items():
                if search_word.lower() in value.lower():
                    print(f"Found '{search_word}' in {key}")
        else:
            search_in_folder(folder_path, search_word)